#coding=utf-8
import base64,hashlib,hmac,json,time
from settings import JWT_SECRET
JWT_HEADER = base64.urlsafe_b64encode('{"typ":"JWT","alg":"HS256"}')
DEFAULT_EXPIRE_TIME=24*3600*15


class JwtGenerator:
    def __init__(self,user_id,expire,session=None):
        self.user_id= user_id
        self.expire = expire
        self.jwt_code = None
        self.session = session

    def get(self):
        if self.jwt_code is None:
            build_time = int(time.time())
            exp_time = build_time+self.expire
            payload = {
                'iat':build_time,
                'iss':'',
                'exp':exp_time,
                'nbf':exp_time,
                'sub':self.user_id,
            }
            if self.session:
                payload['session']=self.session
            payload_str = base64.urlsafe_b64encode(json.dumps(payload)).replace(b'=', b'')
            encode_str = '.'.join([JWT_HEADER, payload_str])
            be_signature = base64.urlsafe_b64encode(
                hmac.new(JWT_SECRET, msg=encode_str, digestmod=hashlib.sha256).digest()).replace(b'=', b'')
            self.jwt_code='.'.join([encode_str,be_signature])
        return self.jwt_code

    @staticmethod
    def vaild(jwt_code):
        header_code, payload_code, signature_code = jwt_code.split('.')
        encode_str = '.'.join([header_code, payload_code])
        be_signature = base64.urlsafe_b64encode(
            hmac.new(JWT_SECRET, msg=encode_str, digestmod=hashlib.sha256).digest()).replace(b'=', b'')
        if signature_code != be_signature:
            return False
        payload_code += (4 - len(payload_code) % 4) * '='
        payload = json.loads(base64.b64decode(payload_code))
        now_time = time.time()
        if payload['exp']<now_time:
            return False
        return JwtGenerator(payload['sub'],None,payload.get('session',{}))

if __name__=='__main__':
    print JwtGenerator(1,24*3600*15).get()
